fix: Add PipelineVariable support to ModelTrainer fields (fixes #5524)#5608
Merged
mollyheamazon merged 2 commits intoaws:masterfrom Mar 9, 2026
Merged
Conversation
) Extend StrPipeVar type to ModelTrainer's direct fields: - training_image: Optional[str] -> Optional[StrPipeVar] - algorithm_name: Optional[str] -> Optional[StrPipeVar] - training_input_mode: Optional[str] -> Optional[StrPipeVar] - environment: Dict[str, str] -> Dict[str, StrPipeVar] This follows the existing V3 pattern already used by SourceCode, OutputDataConfig, and Compute (for instance_type). The StrPipeVar type alias and PipelineVariable.__get_pydantic_core_schema__() already exist in the codebase. This unblocks V2->V3 migration for SageMaker Pipelines users who need to pass ParameterString to ModelTrainer fields. Fixes aws#5524
…ble-safe logging
- Add test_model_trainer_pipeline_variable.py with 9 tests:
- 4 PipelineVariable acceptance tests (training_image, algorithm_name,
training_input_mode, environment)
- 4 regression tests (real string values still work)
- 1 invalid type rejection test
- Fix PipelineVariable-safe logging in model_post_init (avoid __str__
on PipelineVariable which raises TypeError)
All 57 tests pass (48 existing + 9 new, 0 regressions).
mollyheamazon
approved these changes
Mar 9, 2026
Contributor
Author
|
The 3 codestyle-doc-tests failures (sagemaker-train, sagemaker-mlops, sagemaker-serve) are pre-existing on master — confirmed identical failures on unrelated PR #5602. This PR only modifies |
Contributor
|
Thanks for this @amitmodi @mollyheamazon. I could help contribute with the follow up PR? |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR extends PipelineVariable support to ModelTrainer direct fields, unblocking V2-V3 migration for SageMaker Pipelines users.
Fixes: #5524
Problem
ModelTrainer uses Pydantic BaseModel with strict validation. Fields like training_image are typed as Optional[str], causing Pydantic to reject PipelineVariable objects (e.g. ParameterString) with a ValidationError. This blocks any customer using SageMaker Pipelines from migrating to V3.
Solution
Follow the existing V3 pattern already established by SourceCode, OutputDataConfig, and Compute - use the StrPipeVar type alias (Union[str, PipelineVariable]) which already exists in pipeline_variable.py with Pydantic compatibility via get_pydantic_core_schema().
Changes (1 file, 5 insertions, 4 deletions)
Testing
Before: ModelTrainer(training_image=ParameterString(name='img')) raises ValidationError
After: ModelTrainer(training_image=ParameterString(name='img')) works
Note
This is step 1. Integer/boolean fields in nested configs (Compute, Networking) need IntPipeVar/BoolPipeVar in a follow-up PR.